package com.example.wepeaker

import ai.onnxruntime.OnnxTensor
import ai.onnxruntime.OrtEnvironment
import ai.onnxruntime.OrtSession
import android.content.Context
import org.json.JSONArray
import java.io.File
import java.nio.FloatBuffer
import java.util.*
import kotlin.math.*

class Speaker(private val context: Context) {
    private lateinit var onnxSession: OrtSession
    private lateinit var melBasis: Array<FloatArray>
    private lateinit var window: FloatArray

    private val ortEnvironment = OrtEnvironment.getEnvironment()

    init {
        initModel()
    }

    private fun initModel() {
        // Load ONNX model
        val modelFile = File(context.getExternalFilesDir(null), "cnceleb_resnet34.onnx")
        if (!modelFile.exists()) {
            context.assets.open("cnceleb_resnet34.onnx").use { input ->
                modelFile.outputStream().use { output ->
                    input.copyTo(output)
                }
            }
        }

        val sessionOptions = OrtSession.SessionOptions()
        sessionOptions.setInterOpNumThreads(1)
        sessionOptions.setIntraOpNumThreads(1)

        onnxSession = ortEnvironment.createSession(modelFile.absolutePath, sessionOptions)

        // Load window and mel basis from JSON files
        loadResources()
    }

    private fun loadResources() {
        // Load window from JSON
        context.assets.open("wepeaker_window.json").use { inputStream ->
            val jsonString = inputStream.bufferedReader().use { it.readText() }
            val jsonArray = JSONArray(jsonString)
            val windowArray = jsonArray.getJSONArray(0)
            window = FloatArray(windowArray.length())
            for (i in 0 until windowArray.length()) {
                window[i] = windowArray.getDouble(i).toFloat()
            }
        }

        // Load mel basis from JSON
        context.assets.open("wepeaker_mel_basis.json").use { inputStream ->
            val jsonString = inputStream.bufferedReader().use { it.readText() }
            val jsonArray = JSONArray(jsonString)
            melBasis = Array(jsonArray.length()) { FloatArray(jsonArray.getJSONArray(0).length()) }
            for (i in 0 until jsonArray.length()) {
                val row = jsonArray.getJSONArray(i)
                for (j in 0 until row.length()) {
                    melBasis[i][j] = row.getDouble(j).toFloat()
                }
            }
        }
    }

    fun extractEmbeddingFromPcmOnnx(wav: FloatArray): FloatArray {
        println("Input audio check------------------------------------shape:${wav.size},sum:${wav.sum()}")
        val embeddings = FloatArray(256) { 0f }

        var offset = 0
        while (offset < wav.size) {
            val chunkSize = min(16000 * 3, wav.size - offset)
            val wav_ = FloatArray(16000 * 3) { 0f }
            System.arraycopy(wav, offset, wav_, 0, chunkSize)

            val weight = chunkSize.toFloat() / 48000f

            // Multiply by 32768 as in Python code
            for (i in wav_.indices) {
                wav_[i] *= 32768f
            }

            println("Pre-fbank input audio check------------------------------------shape:${wav_.size},sum:${wav_.sum()}")

            // Extract fbank features
            var feats = fbankNumpy(wav_)

            println(
                "Fbank result check------------------------------------shape:${feats.size}x${feats[0].size},sum:${
                feats.fold(0f) { acc, row -> acc + row.sum() }
            }")

            // Mean normalization
            val means = FloatArray(feats[0].size) { j ->
                var sum = 0f
                for (i in feats.indices) {
                    sum += feats[i][j]
                }
                sum / feats.size
            }

            for (i in feats.indices) {
                for (j in feats[i].indices) {
                    feats[i][j] -= means[j]
                }
            }

            // Create ONNX input tensor
            val shape = longArrayOf(1, feats.size.toLong(), feats[0].size.toLong())
            val flatSize = shape[0] * shape[1] * shape[2]
            val flatBuffer = FloatBuffer.allocate(flatSize.toInt())

            for (i in feats.indices) {
                for (j in feats[i].indices) {
                    flatBuffer.put(feats[i][j])
                }
            }
            flatBuffer.rewind()

            println(
                "Feats check------------------------------------shape:${shape.joinToString("x")},sum:${
                    flatBuffer.asReadOnlyBuffer().let {
                        var sum = 0f
                        while (it.hasRemaining()) sum += it.get()
                        sum
                    }
                }"
            )

            // Run inference
            val inputName = "feats"
            val inputTensor = OnnxTensor.createTensor(ortEnvironment, flatBuffer, shape)
            val inputMap = Collections.singletonMap(inputName, inputTensor)
            val output = onnxSession.run(inputMap)

            // Get embeddings and apply weight
            try {
                // Get the output object and extract the tensor data safely
                val outputObj = output.get("embs")

                if (outputObj == null) {
                    println("Error: Output object is null")
                } else {
                    println("Output object type: ${outputObj.javaClass.name}")

                    var tempEmbeddings: FloatArray? = null

                    // Try to extract float array from the output object using reflection 
                    // for maximum compatibility
                    try {
                        if (outputObj.javaClass.name.contains("Optional")) {
                            // Handle Optional type
                            val getMethod = outputObj.javaClass.getMethod("get")
                            val innerObj = getMethod.invoke(outputObj)

                            if (innerObj != null) {
                                println("Optional inner object type: ${innerObj.javaClass.name}")
                                tempEmbeddings = extractEmbeddingsFromObject(innerObj)
                            }
                        } else {
                            // Try direct extraction
                            tempEmbeddings = extractEmbeddingsFromObject(outputObj)
                        }

                        // Apply extracted embeddings
                        if (tempEmbeddings != null) {
                            for (i in embeddings.indices) {
                                embeddings[i] += tempEmbeddings[i] * weight
                            }
                            println("Embeddings check------------------------------------shape:${embeddings.size},sum:${embeddings.sum()}")
                        } else {
                            println("Failed to extract embeddings data from output")
                        }
                    } catch (e: Exception) {
                        println("Error extracting embeddings: ${e.message}")
                        e.printStackTrace()
                    }
                }
            } catch (e: Exception) {
                println("Error processing output object: ${e.message}")
                e.printStackTrace()
            }

            inputTensor.close()

            offset += 16000 * 3
        }

        return embeddings
    }

    fun fbankNumpy(waveform: FloatArray): Array<FloatArray> {
        val preemphasisCoefficient = 0.97f
        val usePower = true
        val windowShift = 160
        val windowSize = 400
        val paddedWindowSize = 512

        // Calculate number of frames
        val numSamples = waveform.size
        val m = 1 + (numSamples - windowSize) / windowShift

        // Create strided input
        val stridedInput = Array(m) { FloatArray(windowSize) }
        for (i in 0 until m) {
            System.arraycopy(waveform, i * windowShift, stridedInput[i], 0, windowSize)
        }

        println(
            "Fbank step1 result check------------------------------------shape:${stridedInput.size}x${stridedInput[0].size},sum:${
            stridedInput.fold(0f) { acc, row -> acc + row.sum() }
        }")

        // Remove mean from each frame
        for (i in stridedInput.indices) {
            val mean = stridedInput[i].average().toFloat()
            for (j in stridedInput[i].indices) {
                stridedInput[i][j] -= mean
            }
        }

        println(
            "Fbank step2 result check------------------------------------shape:${stridedInput.size}x${stridedInput[0].size},sum:${
            stridedInput.fold(0f) { acc, row -> acc + row.sum() }
        }")

        // Preemphasis - first pad the input
        val offsetStridedInput = Array(stridedInput.size) { FloatArray(windowSize + 1) }
        for (i in stridedInput.indices) {
            // Edge padding
            offsetStridedInput[i][0] = stridedInput[i][0]
            System.arraycopy(stridedInput[i], 0, offsetStridedInput[i], 1, windowSize)
        }

        println(
            "Fbank step3 result check------------------------------------shape:${offsetStridedInput.size}x${offsetStridedInput[0].size},sum:${
            offsetStridedInput.fold(0f) { acc, row -> acc + row.sum() }
        }")

        // Apply preemphasis
        for (i in stridedInput.indices) {
            for (j in stridedInput[i].indices) {
                stridedInput[i][j] -= preemphasisCoefficient * offsetStridedInput[i][j]
            }
        }

        println(
            "Fbank step4 result check------------------------------------shape:${stridedInput.size}x${stridedInput[0].size},sum:${
            stridedInput.fold(0f) { acc, row -> acc + row.sum() }
        }")

        // Apply window
        for (i in stridedInput.indices) {
            for (j in stridedInput[i].indices) {
                stridedInput[i][j] *= window[j]
            }
        }

        println(
            "Fbank step5 result check------------------------------------shape:${stridedInput.size}x${stridedInput[0].size},sum:${
            stridedInput.fold(0f) { acc, row -> acc + row.sum() }
        }")

        // Pad to FFT size
        val paddedInput = Array(stridedInput.size) { FloatArray(paddedWindowSize) }
        for (i in stridedInput.indices) {
            System.arraycopy(stridedInput[i], 0, paddedInput[i], 0, windowSize)
            // Remaining part is already zeroed out by FloatArray initialization
        }

        println(
            "Fbank step6 result check------------------------------------shape:${paddedInput.size}x${paddedInput[0].size},sum:${
            paddedInput.fold(0f) { acc, row -> acc + row.sum() }
        }")

        // Compute FFT 
        val spectrum = Array(paddedInput.size) { FloatArray(paddedWindowSize / 2 + 1) }
        for (i in paddedInput.indices) {
            val fft = applyRfft(paddedInput[i])
            for (j in fft.indices) {
                spectrum[i][j] =
                    sqrt(fft[j].real * fft[j].real + fft[j].imag * fft[j].imag).toFloat()
            }
        }

        println(
            "Fbank step7 result check------------------------------------shape:${spectrum.size}x${spectrum[0].size},sum:${
            spectrum.fold(0f) { acc, row -> acc + row.sum() }
        }")

        // Apply power if needed
        if (usePower) {
            for (i in spectrum.indices) {
                for (j in spectrum[i].indices) {
                    spectrum[i][j] = spectrum[i][j] * spectrum[i][j]
                }
            }
        }

        // Apply mel filterbank
        val melEnergies = Array(spectrum.size) { FloatArray(melBasis.size) }
        for (i in spectrum.indices) {
            for (j in melBasis.indices) {
                var sum = 0f
                for (k in spectrum[i].indices) {
                    sum += melBasis[j][k] * spectrum[i][k]
                }
                melEnergies[i][j] = sum
            }
        }

        println(
            "Fbank step8 result check------------------------------------shape:${melEnergies.size}x${melEnergies[0].size},sum:${
            melEnergies.fold(0f) { acc, row -> acc + row.sum() }
        }")

        // Apply log
        for (i in melEnergies.indices) {
            for (j in melEnergies[i].indices) {
                melEnergies[i][j] = ln(max(melEnergies[i][j], 1.1921e-7f))
            }
        }

        return melEnergies
    }

    private fun applyRfft(input: FloatArray): Array<Complex> {
        // Simple implementation of RFFT using basic FFT
        // For production use, you'd want to use a library like JTransforms
        val result = Array(input.size / 2 + 1) { Complex(0.0, 0.0) }

        // Basic implementation of RFFT
        val n = input.size

        // Compute DFT for the first half of frequencies
        for (k in 0 until n / 2 + 1) {
            var sumReal = 0.0
            var sumImag = 0.0

            for (t in 0 until n) {
                val angle = -2.0 * Math.PI * k * t / n
                sumReal += input[t] * cos(angle)
                sumImag += input[t] * sin(angle)
            }

            result[k] = Complex(sumReal, sumImag)
        }

        return result
    }

    fun toRmss(wav: FloatArray, hopSize: Int = 160, windowSize: Int = 1024): FloatArray {
        val fNums = wav.size / hopSize
        val paddedWav = FloatArray(wav.size + windowSize)

        // Reflect padding
        val halfWindow = windowSize / 2
        for (i in 0 until halfWindow) {
            paddedWav[i] = wav[halfWindow - i - 1]
            paddedWav[wav.size + halfWindow + i] = wav[wav.size - i - 1]
        }

        // Copy original array
        System.arraycopy(wav, 0, paddedWav, halfWindow, wav.size)

        val newRmss = FloatArray(fNums)
        for (i in 0 until fNums) {
            var sum = 0f
            for (j in 0 until windowSize) {
                val sample = paddedWav[i * hopSize + j]
                sum += sample * sample
            }
            newRmss[i] = sqrt(sum / windowSize)
        }

        return newRmss
    }

    fun run(wav: FloatArray): Float {
        val rmss = toRmss(wav)
        val setenceEmbeddings = mutableListOf<FloatArray>()
        val setenceRmssWeight = mutableListOf<Float>()

        var i = 0
        while (i < wav.size) {
            val pieceSize = min(16000 * 6, wav.size - i)
            val wavPiece = FloatArray(pieceSize)
            System.arraycopy(wav, i, wavPiece, 0, pieceSize)

            val embedding = extractEmbeddingFromPcmOnnx(wavPiece)

            // Normalize embedding
            val norm = sqrt(embedding.map { it * it }.sum().toDouble()).toFloat()
            val normalizedEmbedding = embedding.map { it / norm }.toFloatArray()

            setenceEmbeddings.add(normalizedEmbedding)
            setenceRmssWeight.add(
                rmss.slice(i / 160 until min(i / 160 + 600, rmss.size)).average().toFloat()
            )

            i += 16000 * 6
        }

        // Convert to array for easier processing
        val setenceEmbeddingsArray = setenceEmbeddings.toTypedArray()

        // Calculate similarity matrix
        val simMatrix =
            Array(setenceEmbeddingsArray.size) { FloatArray(setenceEmbeddingsArray.size) }
        for (i in setenceEmbeddingsArray.indices) {
            for (j in setenceEmbeddingsArray.indices) {
                var dotProduct = 0f
                for (k in setenceEmbeddingsArray[i].indices) {
                    dotProduct += setenceEmbeddingsArray[i][k] * setenceEmbeddingsArray[j][k]
                }
                simMatrix[i][j] = dotProduct
            }
        }

        // Use a more stringent similarity threshold for music content
        val similarityThreshold = 0.75f

        // Find max cluster
        var maxClusterIndex = 0
        var maxCount = 0
        for (i in simMatrix.indices) {
            val count = simMatrix[i].count { it > similarityThreshold }
            if (count > maxCount) {
                maxCount = count
                maxClusterIndex = i
            }
        }

        // Calculate kernel embedding
        var kernelEmbedding = FloatArray(256) { 0f }
        var kernelCount = 0

        for (i in simMatrix[maxClusterIndex].indices) {
            if (simMatrix[maxClusterIndex][i] > similarityThreshold) {
                for (j in setenceEmbeddingsArray[i].indices) {
                    kernelEmbedding[j] += setenceEmbeddingsArray[i][j]
                }
                kernelCount++
            }
        }

        // Normalize kernel embedding
        if (kernelCount > 0) {
            for (i in kernelEmbedding.indices) {
                kernelEmbedding[i] /= kernelCount
            }
        }

        val kernelNorm = sqrt(kernelEmbedding.map { it * it }.sum().toDouble()).toFloat()
        for (i in kernelEmbedding.indices) {
            kernelEmbedding[i] /= kernelNorm
        }

        // Refine kernel embedding
        for (iter in 0 until 5) {
            val kernelDis = FloatArray(setenceEmbeddingsArray.size)
            for (i in setenceEmbeddingsArray.indices) {
                var dotProduct = 0f
                for (j in setenceEmbeddingsArray[i].indices) {
                    dotProduct += kernelEmbedding[j] * setenceEmbeddingsArray[i][j]
                }
                kernelDis[i] = dotProduct
            }

            println("$iter ${kernelDis.count { it > similarityThreshold }.toFloat() / kernelDis.size}")

            // Update kernel embedding
            kernelEmbedding = FloatArray(256) { 0f }
            kernelCount = 0

            for (i in kernelDis.indices) {
                if (kernelDis[i] > similarityThreshold) {
                    for (j in setenceEmbeddingsArray[i].indices) {
                        kernelEmbedding[j] += setenceEmbeddingsArray[i][j]
                    }
                    kernelCount++
                }
            }

            // Normalize kernel embedding
            if (kernelCount > 0) {
                for (i in kernelEmbedding.indices) {
                    kernelEmbedding[i] /= kernelCount
                }
            }

            val newNorm = sqrt(kernelEmbedding.map { it * it }.sum().toDouble()).toFloat()
            for (i in kernelEmbedding.indices) {
                kernelEmbedding[i] /= newNorm
            }
        }

        // Calculate final kernel distance
        val kernelDis = FloatArray(setenceEmbeddingsArray.size)
        for (i in setenceEmbeddingsArray.indices) {
            var dotProduct = 0f
            for (j in setenceEmbeddingsArray[i].indices) {
                dotProduct += kernelEmbedding[j] * setenceEmbeddingsArray[i][j]
            }
            kernelDis[i] = dotProduct
        }

        // Calculate cluster diversity (new metric)
        val clusterSimilarities = mutableListOf<Float>()
        for (i in kernelDis.indices) {
            for (j in i + 1 until kernelDis.size) {
                var dotProduct = 0f
                for (k in setenceEmbeddingsArray[i].indices) {
                    dotProduct += setenceEmbeddingsArray[i][k] * setenceEmbeddingsArray[j][k]
                }
                clusterSimilarities.add(dotProduct)
            }
        }
        
        val diversity = if (clusterSimilarities.isNotEmpty()) {
            1.0f - clusterSimilarities.average().toFloat()
        } else {
            0.0f
        }
        
        // Adjust purity calculation
        val setenceRmssWeightArray = setenceRmssWeight.toFloatArray()
        var purifiedWeight = 0f
        var totalWeight = setenceRmssWeightArray.sum()

        for (i in kernelDis.indices) {
            if (kernelDis[i] > similarityThreshold) {
                purifiedWeight += setenceRmssWeightArray[i]
            }
        }
        
        // Apply diversity factor to reduce purity for diverse audio (like music)
        val purity = purifiedWeight / totalWeight
        
        // For singing detection - reduce purity score if diversity is high
        val adjustedPurity = purity * (1.0f - diversity * 0.5f)
        
        return adjustedPurity
    }

    private fun extractEmbeddingsFromObject(obj: Any): FloatArray? {
        try {
            // Try to get the getValue method that ONNX tensors have
            val getValueMethod = obj.javaClass.getMethod("getValue")
            val valueObj = getValueMethod.invoke(obj)

            if (valueObj is Array<*>) {
                // Try to navigate the array structure to get the FloatArray
                if (valueObj.isNotEmpty() && valueObj[0] is Array<*>) {
                    val inner = valueObj[0] as Array<*>
                    if (inner.isNotEmpty() && inner[0] is FloatArray) {
                        return inner[0] as FloatArray
                    }
                } else if (valueObj.isNotEmpty() && valueObj[0] is FloatArray) {
                    return valueObj[0] as FloatArray
                }
            }
        } catch (e: Exception) {
            println("Failed to extract embeddings: ${e.message}")
        }
        return null
    }

    // Helper class for complex numbers
    data class Complex(val real: Double, val imag: Double)

    // Function to release resources
    fun close() {
        if (::onnxSession.isInitialized) {
            onnxSession.close()
        }
        ortEnvironment.close()
    }
}